import sys
sys.path.append('.')
from common_imports import *

# --- 1. Data and Parameters ---
# Data
thresholds = np.array([0, .1, .25, .5, 1, 1.5, 2, 2.5, 3])
means_sd = np.array([0.7001, 0.7093, 0.6989, 0.6733, 0.5975, 0.5299, 0.4763, 0.4319, 0.3853])
means_dr = np.array([0.8173, 0.7248, 0.5737, 0.3783, 0.3, 0.2473, 0.1995, 0.1659, 0.1425])
fraction_sd = np.array([0.7663, 0.4159, 0.3034, 0.1927, 0.0926, 0.0535, 0.0353, 0.0253, 0.0188])
fraction_dr = np.array([0.7663, 0.1873, 0.0624, 0.0245, 0.013, 0.0083, 0.0059, 0.0045, 0.0037])

# A professional, colorblind-friendly palette (teal and orange)
colors = ['#1b9e77', '#d95f02'] # Teal for Semidiurnal, Orange for Diurnal

# Slice data to the desired threshold
max_thresh_idx = np.where(thresholds <= 3)[0][-1] + 1
x = thresholds[:max_thresh_idx]
y1_sd_mean = means_sd[:max_thresh_idx]
y2_dr_mean = means_dr[:max_thresh_idx]
a1_sd_frac = fraction_sd[:max_thresh_idx]
a2_dr_frac = fraction_dr[:max_thresh_idx]


# --- 2. Plotting ---
fig_width_cm = 9
fig_width_in = fig_width_cm / 2.54
fig_height_max_cm = 10
fig_height_in = fig_height_max_cm / 2.54 

fig, ax = plt.subplots(figsize=(fig_width_in, fig_height_in),
                        constrained_layout=True)

# Plot mean incoherent fraction (solid lines)
ax.plot(x, y1_sd_mean, color=colors[0], linestyle='-', marker='o', linewidth=1, markersize=2,
        label='Semidiurnal: Mean Incoherent Fraction')
ax.plot(x, y2_dr_mean, color=colors[1], linestyle='-', marker='o', linewidth=1, markersize=2,
        label='Diurnal: Mean Incoherent Fraction')

# Plot fraction of grid points (dashed lines)
ax.plot(x, a1_sd_frac, color=colors[0], linestyle='--', marker='^', markersize=2, linewidth=1,
        label='Semidiurnal: Fraction of Grid Points')
ax.plot(x, a2_dr_frac, color=colors[1], linestyle='--', marker='^', markersize=2, linewidth=1,
        label='Diurnal: Fraction of Grid Points')


# --- 3. Aesthetics and Labels ---
ax.set_xlabel('Bandpass Temporal Variance Threshold (cm$^2$)', fontsize=9)
ax.set_ylabel('Fraction', fontsize=9)
# ax.set_title('Incoherent Internal Tide Fraction vs. Variance Threshold', fontsize=9, pad=15)
ax.set_xlim(left=-0.16, right=3.1)
ax.set_ylim(bottom=0, top=1.)
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.spines['left'].set_linewidth(1.3)
ax.spines['bottom'].set_linewidth(1.3)
# ax.grid(True, linestyle=':', linewidth=0.6, color='grey')
ax.axvline(x=0, color='grey', linestyle='-', linewidth=0.8, alpha=0.7)
ax.text(-0.1, 0.3, "Deep Ocean without strong eddies", rotation=90, va='center', ha='left', fontsize=7, color='grey')

# --- 4. Annotations ---
# This section adds text labels with arrows pointing to specific data points
# We use ax.annotate for precise control

# Annotations at the first threshold (x=0.1)
idx = 1
ax.annotate(f'{y2_dr_mean[idx]:.2f}', xy=(x[idx], y2_dr_mean[idx]),
            xytext=(x[idx] + 0, y2_dr_mean[idx] +.1), color=colors[1],
            fontweight='bold', arrowprops=dict(arrowstyle='->', color='gray'), fontsize = 7)

idx = 2
ax.annotate(f'{y1_sd_mean[idx]:.2f}', xy=(x[idx], y1_sd_mean[idx]),
            xytext=(x[idx] + 0.1, y1_sd_mean[idx] + 0.05), color=colors[0],
            fontweight='bold', arrowprops=dict(arrowstyle='->', color='gray'), fontsize = 7)
ax.annotate(f'{y2_dr_mean[idx]:.2f}', xy=(x[idx], y2_dr_mean[idx]),
            xytext=(x[idx] + 0.15, y2_dr_mean[idx] - 0.05), color=colors[1],
            fontweight='bold', arrowprops=dict(arrowstyle='->', color='gray'), fontsize = 7)

# Annotations at the 1.0 cm^2 threshold
idx = 4
ax.annotate(f'{y1_sd_mean[idx]:.2f}', xy=(x[idx], y1_sd_mean[idx]),
            xytext=(x[idx] - 0., y1_sd_mean[idx] + 0.1), color=colors[0],
            fontweight='bold', arrowprops=dict(arrowstyle='->', color='gray'), fontsize = 7)
ax.annotate(f'{y2_dr_mean[idx]:.2f}', xy=(x[idx], y2_dr_mean[idx]),
            xytext=(x[idx] + 0.1, y2_dr_mean[idx] - 0.1), color=colors[1],
            fontweight='bold', arrowprops=dict(arrowstyle='->', color='gray'), fontsize = 7)

# Annotations at the 3.0 cm^2 threshold
idx = 8
ax.annotate(f'{y1_sd_mean[idx]:.2f}', xy=(x[idx], y1_sd_mean[idx]),
            xytext=(x[idx] - 0.5, y1_sd_mean[idx] + 0.1), color=colors[0],
            fontweight='bold', arrowprops=dict(arrowstyle='->', color='gray'), fontsize = 7)
ax.annotate(f'{y2_dr_mean[idx]:.2f}', xy=(x[idx], y2_dr_mean[idx]),
            xytext=(x[idx] - 0.5, y2_dr_mean[idx] - 0.05), color=colors[1],
            fontweight='bold', arrowprops=dict(arrowstyle='->', color='gray'), fontsize = 7)


# --- 5. Legend and Layout ---
ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.14), frameon=True, fontsize=9, ncol=1)
ax.tick_params(axis='both', which='major', length=4, width=1.2, labelsize=8)
plt.savefig('Fig6_threshold_dependence_incoherent_fraction.pdf', bbox_inches='tight')